module Graph where

import Data.Vector as V hiding((++),filter, map, foldl)
import Prelude hiding(length)
import Data.Maybe

newtype Graph a = G (Vector [(Int,a)])
    deriving(Show)

-- Возвращает пустой граф
empty :: Graph a
empty = G V.empty

fromEdges :: [(Int,Int,a)] -> Graph a
fromEdges = foldl resizeAndAdd Graph.empty
    where resizeAndAdd g (i,j,w) 
              | i < 1 || j < 1 = error "Wrong index"
              | otherwise = addEdge (resize g $ max i j) i j w
          resize g s 
              | s > size g = foldl (\g' _ -> addVertex g') g [j|j<-[size g + 1 .. s]]
              | otherwise  = g

isolatedVertixies :: Int -> Graph a
isolatedVertixies n 
        | n < 0 = error $ "Can't build graph with " ++ show n ++ " vertixies"
        | otherwise = repeatAdd n Graph.empty
    where repeatAdd 0 g = g
          repeatAdd n g = repeatAdd (n-1) (addVertex g)

size :: Graph a -> Int
size (G v) = length v

-- Добавляет вершину в граф
addVertex :: Graph a -> Graph a
addVertex (G v)  = G $ snoc v []

-- Удаляет вершину с заданным номером из графа
delVertex :: Graph a -> Int -> Graph a
delVertex (G v) i 
        | i > length v || i < 1 = error $ "Vertex index " ++ show i ++ " out of bounds."
        | otherwise = G $ generate (length v - 1) genVertexies 
    where ind = i - 1
          genVertexies ind' = 
              map (\e@(i',d) -> if i' < i then e else (i'-1,d)) $  -- сдвигаем индексы поздних вершин
                  filter (\(i',_) -> i' /= i) -- удаляем все рёбра в эту вершину
                         (v ! (if ind' >= ind then ind' + 1 else ind') )

-- Добавляет ребро из @arg2 -> @arg3
addEdge :: Graph a -> Int -> Int -> a -> Graph a
addEdge (G v) i1 i2 d 
        | i1 < 1 || i2 < 1 || i1 > length v || i2 > length v = error $ "Vertex index " ++ show i1 ++ ", or " ++ show i2 ++ " out of bounds."
        | otherwise = G $ v // [(i1 - 1, (i2,d):(v ! (i1 - 1)))] -- добавляем к i1 новое ребро до i2 с весом d

-- Добавляет ребро из @arg2 -> @arg3 и наоборот
addBidirectEdge :: Graph a -> Int -> Int -> a -> Graph a
addBidirectEdge (G v) i1 i2 d 
        | i1 < 1 || i2 < 1 || i1 > length v || i2 > length v = error $ "Vertex index " ++ show i1 ++ ", or " ++ show i2 ++ " out of bounds."
        | otherwise = G $ v // [(i1 - 1, (i2,d):(v ! (i1 - 1))),
                                (i2 - 1,(i1,d):(v ! (i2 - 1)))] -- добавляем к i1 новое ребро до i2 с весом d и наоборот


-- Добавляет ребро из @arg2 -> @arg3, удаляя все предыдущие
updateEdge :: Graph a -> Int -> Int -> a -> Graph a
updateEdge (G v) i1 i2 d 
        | i1 < 1 || i2 < 1 || i1 > length v || i2 > length v = error $ "Vertex index " ++ show i1 ++ ", or " ++ show i2 ++ " out of bounds."
        | otherwise = G $ v // [(i1 - 1, (i2,d):(filter (\(i,_) -> i /= i2) $ v ! (i1 - 1)))] -- добавляем к i1 новое ребро до i2 с весом d, а из предыдущих удаляем


-- Добавляет ребро из @arg2 -> @arg3 и обратно, удаляя все предыдущие
updateBidirectEdge :: Graph a -> Int -> Int -> a -> Graph a
updateBidirectEdge (G v) i1 i2 d 
        | i1 < 1 || i2 < 1 || i1 > length v || i2 > length v = error $ "Vertex index " ++ show i1 ++ ", or " ++ show i2 ++ " out of bounds."
        | otherwise = G $ v // [(i1 - 1, (i2,d):(filter (\(i,_) -> i /= i2) $ v ! (i1 - 1))), -- добавляем к i1 новое ребро до i2 с весом d, а из предыдущих удаляем
                                (i2 - 1, (i1,d):(filter (\(i,_) -> i /= i1) $ v ! (i2 - 1)))] -- добавляем к i2 новое ребро до i1 с весом d, а из предыдущих удаляем

-- Удаляет ребро из @arg2 -> @arg3 с заданым весом
delWeightedEdge :: Eq a => Graph a -> Int -> Int -> a -> Graph a
delWeightedEdge (G v) i1 i2 d 
        | i1 < 1 || i2 < 1 || i1 > length v || i2 > length v = error $ "Vertex index " ++ show i1 ++ ", or " ++ show i2 ++ " out of bounds."
        | otherwise = G $ v // [(i1 - 1, filter (\(i,d') -> i /= i2 || d' /= d) $ v ! (i1 - 1))] -- удаляем ребро если к вершине i2 и с весом d

-- Удаляет ребро из @arg2 -> @arg3
delEdges :: Graph a -> Int -> Int -> Graph a
delEdges (G v) i1 i2 
        | i1 < 1 || i2 < 1 || i1 > length v || i2 > length v = error $ "Vertex index " ++ show i1 ++ ", or " ++ show i2 ++ " out of bounds."
        | otherwise = G $ v // [(i1 - 1, filter (\(i,_) -> i /= i2) $ v ! (i1 - 1))] -- удаляем ребро если к вершине i2

-- Удаляет ребро из @arg2 -> @arg3 и обратное
delBidirectEdges :: Graph a -> Int -> Int -> Graph a
delBidirectEdges (G v) i1 i2 
        | i1 < 1 || i2 < 1 || i1 > length v || i2 > length v = error $ "Vertex index " ++ show i1 ++ ", or " ++ show i2 ++ " out of bounds."
        | otherwise = G $ v // [(i1 - 1, filter (\(i,_) -> i /= i2) $ v ! (i1 - 1)), -- удаляем ребро если к вершине i2
                                (i2 - 1, filter (\(i,_) -> i /= i1) $ v ! (i2 - 1))] -- удаляем ребро если к вершине i1

areAdjacent :: Graph a -> Int -> Int -> Bool
areAdjacent (G v) i1 i2 
        | i1 < 1 || i2 < 1 || i1 > length v || i2 > length v = error $ "Vertex index " ++ show i1 ++ ", or " ++ show i2 ++ " out of bounds."
        | otherwise = isJust $ lookup i2 $ v ! (i1 - 1) -- возвращаем True если lookup нашел ключ

getAdjacents :: Graph a -> Int -> [Int]
getAdjacents (G v) i
        | i < 1 || i > length v = error $ "Vertex index " ++ show i ++ " out of bounds."
        | otherwise = filterDublicates $ map fst $ v ! (i - 1)
    where filterDublicates [] = []
          filterDublicates (i:is) = i:(filter (\i' -> i' /= i) is)

getWeight :: Graph a -> Int -> Int -> Maybe a
getWeight (G v) i1 i2 
        | i1 < 1 || i2 < 1 || i1 > length v || i2 > length v = error $ "Vertex index " ++ show i1 ++ ", or " ++ show i2 ++ " out of bounds."
        | otherwise = lookup i2 $ v ! (i1 - 1)

getWeights :: Graph a -> Int -> Int -> [a]
getWeights (G v) i1 i2 
        | i1 < 1 || i2 < 1 || i1 > length v || i2 > length v = error $ "Vertex index " ++ show i1 ++ ", or " ++ show i2 ++ " out of bounds."
        | otherwise = map snd $ filter (\(i,_) -> i == i2) $ v ! (i1 - 1)

mapEdges :: Graph a -> ((Int, Int, a) -> Maybe b) -> Graph b
mapEdges g f = 
    foldl (\g' i -> 
              foldl (\g'' j -> 
                        case f (i,j,fromJust $ getWeight g i j) of
                         Nothing -> g''
                         Just b  -> addEdge g'' i j b) 
                    g' $ getAdjacents g' i) 
          Graph.empty [i|i <- [1..size g]]

type SimpleGraph = Graph Bool

addSimpleGraphEdge :: SimpleGraph -> Int -> Int -> SimpleGraph
addSimpleGraphEdge g i1 i2 = addEdge (addEdge g i1 i2 False) i2 i1 False


